"""
PostgreSQL Sink 函数
用于 Flink DataStream 写入 PostgreSQL
"""
import os
import json
import logging
import psycopg2
from typing import Any

logger = logging.getLogger(__name__)


def write_to_postgres(value: Any) -> None:
    """
    将 Flink Agent 输出写入 PostgreSQL

    Args:
        value: Agent 输出的数据（可能是 dict、str 或 pyflink.common.Row）
    """
    try:
        # 处理不同的输入格式
        if hasattr(value, 'as_dict'):
            # pyflink.common.Row 对象
            data = value.as_dict()
        elif isinstance(value, str):
            # JSON 字符串
            data = json.loads(value)
        elif isinstance(value, dict):
            # 字典
            data = value
        else:
            logger.warning(f"未知数据格式: {type(value)}")
            return

        # 如果是包装格式，提取实际数据
        if 'user_id' not in data and len(data) == 1:
            data = list(data.values())[0]

        # 验证必需字段
        required_fields = ['user_id', 'url', 'behavior_type', 'behavior_description']
        for field in required_fields:
            if field not in data:
                logger.error(f"缺少必需字段: {field}, 数据: {data}")
                return

        # 连接数据库
        conn = psycopg2.connect(
            host=os.getenv('POSTGRES_HOST', 'postgres'),
            port=int(os.getenv('POSTGRES_PORT', 5432)),
            database=os.getenv('POSTGRES_DB', 'behavior_analysis'),
            user=os.getenv('POSTGRES_USER', 'postgres'),
            password=os.getenv('POSTGRES_PASSWORD', 'postgres')
        )
        cursor = conn.cursor()

        # 写入数据
        insert_query = """
            INSERT INTO user_behaviors (
                user_id, url, behavior_type, behavior_description,
                confidence, analysis_method, created_at
            ) VALUES (%s, %s, %s, %s, %s, %s, NOW())
        """

        cursor.execute(insert_query, (
            data.get('user_id'),
            data.get('url'),
            data.get('behavior_type'),
            data.get('behavior_description'),
            data.get('confidence', 0.8),
            data.get('analysis_method', 'flink_agents_llm')
        ))

        conn.commit()
        cursor.close()
        conn.close()

        logger.debug(f"✅ 写入 PostgreSQL: {data.get('user_id')} - {data.get('behavior_type')}")

    except Exception as e:
        logger.error(f"❌ 写入 PostgreSQL 失败: {e}")
        import traceback
        traceback.print_exc()
